import numpy as np
import os


_url_base = '/Users/rench/code/data/MNIST/ubyte'
_filename = {
    'train_imgs' : 'train-images-idx3-ubyte',
    'train_labs' : 'train-labels-idx1-ubyte',
    'test_imgs' : 't10k-images-idx3-ubyte',
    'test_labs' : 't10k-labels-idx1-ubyte'
}


class MNIST(object):
    def __init__(self, root=_url_base):
        self.root = root
        self.filename = _filename

    def train_sets(self):
        return self._read_file('train_imgs').reshape(-1,28,28),self._read_file('train_labs',lable=True)

    def test_sets(self):
        return self._read_file('test_imgs').reshape(-1,28,28),self._read_file('test_labs',lable=True)
    
    def _read_file(self, name, lable=False):
        if lable:
            with open(os.path.join(self.root, self.filename[name]),"rb") as f:
                file = f.read()
                data = np.frombuffer(file, np.uint8, offset=8).reshape(-1,1)

            return data
        else:
            with open(os.path.join(self.root, self.filename[name]),"rb") as f:
                file = f.read()
                data = np.frombuffer(file, np.uint8, offset=16).reshape(-1,28,28) 
            return data          

    def get_sets(self,train=True):
        if(train):
            return self.train_sets()
        else:
            return self.test_sets()
